import cv2
import numpy as np
from sklearn import svm
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score


# 1. 图像预处理
def preprocess_image(image_path):
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    return blurred


# 2. 连通域分割
def connected_component_segmentation(image):
    # 边缘检测
    edges = cv2.Canny(image, 50, 150)

    # 连通域标记
    num_labels, labels_im = cv2.connectedComponents(edges)

    return num_labels, labels_im


# 3. 特征提取
def extract_features(labels_im, image):
    features = []
    for label in range(1, np.max(labels_im) + 1):
        mask = np.zeros_like(image)
        mask[labels_im == label] = 255  # 创建连通区域的掩膜

        # 提取特征，简单例子使用区域面积和中心坐标
        area = np.sum(mask > 0)
        moments = cv2.moments(mask)
        if moments["m00"] != 0:
            cX = int(moments["m10"] / moments["m00"])
            cY = int(moments["m01"] / moments["m00"])
            features.append([area, cX, cY])
    return features


# 4. SVM训练和预测
def train_svm(features, labels):
    features = np.array(features)
    labels = np.array(labels)

    X_train, X_test, y_train, y_test = train_test_split(features, labels, test_size=0.2, random_state=42)
    clf = svm.SVC(kernel='linear')
    clf.fit(X_train, y_train)

    # 预测和评估
    y_pred = clf.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f"SVM Accuracy: {accuracy:.2f}")

    return clf


# 主函数
def main(image_path, labels):
    # 进行图像预处理
    preprocessed_image = preprocess_image(image_path)

    # 进行连通域分割
    num_labels, labels_im = connected_component_segmentation(preprocessed_image)

    # 提取特征
    features = extract_features(labels_im, preprocessed_image)

    # 训练SVM模型
    svm_model = train_svm(features, labels)

    # 返回训练好的模型
    return svm_model


# 测试代码
if __name__ == "__main__":
    # 假设我们有示例图像和其对应的标签
    image_path = "city_railway_image.jpg"  # 请替换为您的图像路径
    labels = [0, 1, 0, 1, 0]  # 示例标签，替换为实际标签

    trained_model = main(image_path, labels)